const fs = require('fs').promises;
const path = require('path');
const BaseDatasetWrapper = require('../../BaseDatasetWrapper');
const YOLOLabelLoader = require('./YOLOLabelLoader');
const { DatasetListDB } = require('../../../dbs/dbDatasetList');
const { getAllFilesInDirectory } = require('../../../utils/fileUtils');

class YOLODatasetWrapper extends BaseDatasetWrapper {
    // 读入目录信息，解析对应的标注结果
    async parse(dataset_info) {
        // console.log('YOLODatasetWrapper::parse dataset_info = ', dataset_info);
        const datasetId = dataset_info.dataset_id;
        const datasetDir = dataset_info.dataset_dir;
        this.parseCallback(datasetId, 'parse_status', '读取图片列表...');
        let imageFiles = await getAllFilesInDirectory(datasetDir, ['.jpg', '.jpeg', '.png', '.bmp'], ['workspace']);
        this.parseCallback(datasetId, 'parse_status', '读取标注文件列表...');
        let labelFilePaths = await getAllFilesInDirectory(datasetDir, ['.txt'], ['workspace']);
        this.parseCallback(datasetId, 'parse_start', imageFiles.length);

        // 创建映射关系
        let labelMap = {};
        let categories_name = new Set();
        await Promise.all(
            labelFilePaths.map(async (labelPath) => {
                const labelFileName = path.basename(labelPath);
                labelMap[labelFileName] = labelPath;

                let categories = await this.getCategories(labelPath);
                categories.forEach((name) => categories_name.add(name));
            })
        );

        // 创建工作目录用于存放复制的标注文件
        const workDir = path.join(datasetDir, 'workspace', datasetId);
        await fs.mkdir(workDir, { recursive: true });

        // 加载每一张图片对应的标注数据
        let totalSamples = imageFiles.length;
        let labeledSamples = 0;
        let matchingPairs = [];
        for (let i = 0; i < imageFiles.length; i++) {
            let imageFile = imageFiles[i];
            const labelFileName = path.basename(imageFile, path.extname(imageFile)) + '.txt';
            const labelFilePath = labelMap[labelFileName];

            // 获取图片相对于 dataset_dir 的相对路径
            const imageRelativePath = path.relative(datasetDir, imageFile);
            const labelRelativeDir = path.dirname(imageRelativePath);
            const newLabelFileDir = path.join(workDir, labelRelativeDir);
            const newLabelFilePath = path.join(newLabelFileDir, labelFileName);

            // 确保目标目录存在
            await fs.mkdir(newLabelFileDir, { recursive: true });

            // 拷贝或新建标注文件
            if (labelFilePath) {
                await fs.copyFile(labelFilePath, newLabelFilePath);
                labeledSamples++;
            } else {
                await fs.writeFile(newLabelFilePath, '', 'utf8');
            }

            // 保存图像和标注的配对文件
            matchingPairs.push({
                image_path: imageRelativePath,
                label_path: path.relative(datasetDir, newLabelFilePath)
            });

            // 更新数据加载进展
            this.parseCallback(datasetId, 'parse_percent', (((i + 1) * 100) / totalSamples).toFixed(2));
        }

        // 保存matching pairs结果
        await this.saveMatchingPairs(datasetId, matchingPairs);

        // 汇总图片目录名称，方便后续切分训练集和验证集使用
        let imageDirs = new Set();
        imageFiles.forEach((file) => imageDirs.add(path.dirname(file)));
        await this.createSubsets(dataset_info, Array.from(imageDirs));

        // 保存categories信息
        let categories = Array.from(categories_name).map((category, index) => {
            return {
                id: category,
                name: 'unknown',
                supercategory: 'unknown'
            };
        });
        // console.log(categories_name);
        // console.log(categories);
        this.saveCategories(dataset_info.dataset_id, categories);

        return {
            code: 0,
            msg: {
                totalSamples: totalSamples,
                labeledSamples: labeledSamples
            }
        };
    }

    // 从yolo的txt文件中，读取category索引编号
    async getCategories(label_path) {
        try {
            const txtData = await fs.readFile(label_path, 'utf8');
            const lines = txtData.trim().split('\n');
            let categories = {};
            for (let i = 0; i < lines.length; i++) {
                const parts = lines[i].split(' ').map((x) => parseInt(x));
                if (parts.length !== 5) {
                    // console.warn(`Invalid line format in file: ${label_path}, line: ${i}`);
                    continue;
                }

                let class_id = parts[0];
                if (class_id !== undefined && class_id !== null && class_id !== '') {
                    categories[class_id] = 1;
                }
            }
            return Object.keys(categories);
        } catch (error) {
            console.error(`Error load yolo label file ${label_path},  error: `, error);
            return [];
        }
    }

    // 导出数据集标注结果列表
    async export(toFormat, dataset_info, subsets_info) {
        console.log('YOLODatasetWrapper::export dataset_info = ', dataset_info);
        console.log('YOLODatasetWrapper::export subsets_info = ', subsets_info);

        // 需要将yolo的类别保存到classes.names/obj.names文件中
    }

    async convert(toFormat, dataset_info) {
        let format = toFormat.toLowerCase();
        if (format === 'yolo') {
            return true;
        }
        const datasetId = dataset_info.id;
        const datasetDir = dataset_info.dataset_dir;

        try {
            // 获取数据集的标注列表
            let dblist = this.app.list_dbs[datasetId];
            if (!dblist) {
                this.app.list_dbs[datasetId] = new DatasetListDB(this.app, datasetId);
                dblist = this.app.list_dbs[datasetId];
            }
            let pairs = await dblist.pairs.getPairs();

            // 动态加载labelloader类
            const srcLoader = new YOLOLabelLoader(this.app);
            const loaderFile = toFormat.toUpperCase() + 'LabelLoader.js';
            const loaderClass = require(path.join(__dirname, loaderFile));
            const dstloader = new loaderClass(this.app);

            // 对于每一个pair进行转换
            let categories = await this.app.db.categories.getCategories(datasetId);
            for (let i = 0; i < pairs.length; i++) {
                let pair = pairs[i];
                pair.image_path = path.join(datasetDir, pair.image_path);
                pair.label_path = path.join(datasetDir, pair.label_path);
                let label_data = await srcLoader.load({
                    dataset_id: datasetId,
                    image_path: pair.image_path,
                    label_path: pair.label_path
                });
                // console.log(`label_data = `, label_data);
                let label_wrapper = this.wrapLabelData(dataset_info, pair, label_data);
                let label_info = await dstloader.toString(label_wrapper, categories);

                // 把结果保存到本地文件中
                let outputPath = BaseDatasetWrapper.getConvertedPath(toFormat, dataset_info, pair);
                await fs.writeFile(outputPath, label_info);
            }
            return { code: 0, msg: `YOLO to ${toFormat.toUpperCase()} successfully` };
        } catch (e) {
            console.error('YOLODatasetWrapper::convert error = ', e);
            return { code: -1, msg: e.message };
        }
    }
}

module.exports = YOLODatasetWrapper;
